#include "haydockmethod.h"


HaydockMethod::HaydockMethod() : 
	// A pointer to the used system. The system contains variables 
	// needed to construct elements (coupling strengths etc.).
	system(&System::Std),
	// A double containing the current bias (for bi-layered 
	// systems).
	bias(0.),
	// A double vector containing the current k-point.
	k(arma::zeros<vec>(3)),
	// A  double vector containing the energy sample points used in
	// the calculations. The size of this has little influence on 
	// the operation time.
	E(arma::zeros<vec>(1)),
	// A double containing the broadening factor.
	eta(0.001),
	// An integer containing the number of lanczos states to use in the
	// calculations of the tri-diagonal matrix H_tilde.
	N_states(1000),
	// An integer containing the number of lanczos states to use in the
	// calculations of the tri-diagonal matrix H_tilde.
	N_rand(1)
{ }

// This operator (use by simply appending () to a HaydockMethod
// type variable) calculates the Green's function element for the
// state X. The operator returns the object it self for easy access
// to result. The result is stored in the G result variable. The 
// variable may be retrieved immidiately after calculation i.e:
// hm(id,X).G // Calculate and retreive result.
HaydockMethod & HaydockMethod::operator()(int id, cx_vec X)
{
	// First we retrieve the Hamilton!
	sp_cx_mat H = getHamiltonMatrix(id);
	
	// We then allocate the result.
	G.set_size(E.n_elem);
	
	// A ratio is set according to how challenging the calc. are!
	static const double lanczos_to_haydock_ratio = 0.8;
	// A modified updateProgress is made (as this method should only go to ratio)!
	static const auto updateLaczosProgress = [&] (const double &progress_step) {
		updateProgress(progress_step*lanczos_to_haydock_ratio);
	};
	
	// We find the tri-diagonal elements by using the Lanczos
	// method.
	LanczosMethod lm;
	lm.N_states = N_states;
	if(updateProgress)
		lm.updateProgress = updateLaczosProgress;

	lm(H,X);
	
	// The update step is calculated! Note that updateProgress(step)
	// must be called N times = (1 - ratio)/step. This will give a full 
	// update of 1!
	step = (1. - lanczos_to_haydock_ratio)/(double)(E.n_elem);

	// Then we find the greens value for each energy!
	_Cilk_for(long nE = 0; nE < E.n_elem; nE++)
	{
		G(nE) = calcGreens(lm.alpha, lm.beta, complex(E(nE),eta));

		try_updateProgress();
	}

	//Done!
	return *this;
}

// This operator (use by simply appending () to a HaydockMethod
// type variable) calculates the avg.  Green's function element. 
// The operator returns the object it self for easy access	to 
// result. The result is stored in the G result variable. The 
// variable may be retrieved immidiately after calculation i.e:
// hm(id).G // Calculate and retreive result.
HaydockMethod & HaydockMethod::operator()(int id)
{
	// First we retrieve the Hamilton!
	sp_cx_mat H = getHamiltonMatrix(id);
	
	// We then allocate the result.
	cx_mat fullG(E.n_elem,N_rand);

	// We need to take an avg. across N_rand.
	_Cilk_for(long r = 0; r < N_rand; r++)
	{
		// We generate a random phase.
		cx_vec X = randomPhaseState(H.n_cols);
		
		// We find the tri-diagonal elements by using the Lanczos
		// method.
		LanczosMethod lm;
		lm.N_states = N_states;

		lm(H,X);

		// For each random state we find the greens value for each energy!
		for(long nE = 0; nE < E.n_elem; nE++)
		{
			// We set a weighted Greens value to sum the avg.
			fullG(nE,r) = calcGreens(lm.alpha, lm.beta, complex(E(nE),eta)) / (double)N_rand;
		}
	}

	// Then we take the sum (which is the avg.)
	G = arma::sum(fullG,1);

	//Done!
	return *this;
}


// This function attempts to update the progress.
void HaydockMethod::try_updateProgress() const
{
	if(!updateProgress)
		return;

	updateProgress(step);
}

// This function returns the Hamilton matrix H.
sp_cx_mat HaydockMethod::getHamiltonMatrix(int id) const
{
	// Make a Hamiton object and set the options.
	Hamiltonian hamilton;
	hamilton.system = system;
	hamilton.bias = bias;
	hamilton.k = k;

	// Calculate and return the sparse matrix.
	return hamilton(id).H;
}

// This function returns a random phase state, i.e |r> = 1/R sum_1^R exp(i phi) |i>.
cx_vec HaydockMethod::randomPhaseState(size_t n)
{
	// constants
	static const complex i = complex(0.,1.);
	static const complex pi = arma::datum::pi;
	
	cx_vec phases = 2.*pi*i*arma::randu<vec>(n); // Random phases!
	return arma::exp(phases)/sqrt((double)n); // return normalized!
}

// This function returns the Greens matrix G.
complex HaydockMethod::calcGreens(const vec &alpha, const vec &beta, const complex &z)
{
	const long last = alpha.n_rows - 1;

	// We then find self-consistantly the d-level greens element (the current green element)!
	complex Gc = (z - alpha(last) - sqrt(pow(z-alpha(last),2) - 4. * pow(beta(last-1),2) ) )/( 2. * pow(beta(last-1),2) );

	if(Gc.imag() >= 0.) // To insure a retarded Green's function.
		Gc = (z - alpha(last) + sqrt(pow(z-alpha(last),2) - 4. * pow(beta(last-1),2) ) )/( 2. * pow(beta(last-1),2) );

	// Now we recursively do the continueous fraction.
	for(long i = last-1; i >= 0; i--)
		Gc = 1./(z - alpha(i) - pow(beta(i),2)*Gc);

	// Then return the result!
	return Gc;
}